from pymongo import MongoClient

from core.config import config, BASE_DIR

connection_string = f'mongodb://{config.database.username}:{config.database.password}@{config.database.host}:{config.database.port}/'
client = MongoClient(connection_string)
db = client[config.database.database]
upload_records = db['upload_records']
rbs = db['rbs']
reviews = db['reviews']

# 数据库保存上传记录
def save_upload_record(upload_id, file_name, upload_time, has_return_data, folder_path):
    upload_records.insert_one({
        "upload_id": upload_id,
        "file_name": file_name,
        "upload_time": upload_time,
        "has_return_data": has_return_data,
        "folder_path": str(folder_path) # 路径原先的类型为Path不是str
    })

# 数据库保存所有数据结果，如果rbId已经存在更新数据
def save_prediction_results(upload_id, upload_time, result_data):
    for index, row in result_data.iterrows():
        rbId = row['CLLI_OID']
        # 查询条件，根据rbId进行查询
        query = {"rbId": rbId}
        # 要更新的数据，不包含 createtime
        new_data = {
            "updateTime": upload_time,
            "state": int(row['predictions']), # 根据诈骗情况直接放入审核状态，如果为1则直接放入审核状态
            "fraud": int(row['predictions']),
            "score": float(row['probability']),
            "level": int(row['pre_level']),
            "upload_id": upload_id,
            "APP_AMT": float(row['APP_AMT']),  
            "MBR_NO": row['MBR_NO'],      
            "POHO_NO": row['POHO_NO']    
        }
        # 当插入新文档时设置createtime
        insert_data = {
            "createTime": upload_time
        }
        # 使用 update_one 方法，upsert=True 表示如果不存在则插入，存在则更新
        rbs.update_one(
            query,
            {
                "$set": new_data,
                "$setOnInsert": insert_data
            },
            upsert=True
        )

# 将中高风险的数据插入到审核表当中
def save_reviews(upload_id, result_data):
    for index, row in result_data.iterrows():
        rbId = row['CLLI_OID']
               # 查询条件，根据rbId进行查询
        query = {"rbId": rbId}
        if(int(row['predictions']) == 1):
            new_data = {
                "fraud": int(row['predictions']),
                "score": float(row['probability']),
                "level": int(row['pre_level']),
                "upload_id": upload_id,
                "onlineReview": 0,
                "offlineReview": 0,
                "remark": "",
                "reviewFraud": 2,
            }
            # 使用 update_one 方法，upsert=True 表示如果不存在则插入，存在则更新
            reviews.update_one(
                query,
                {
                    "$set": new_data
                },
                upsert=True
            )